import yaml
import os
import pickle
import numpy as np
import sys
import shutil
import csv


def read_yml(file):
    """读取yml,传入文件路径file"""
    f = open(file, 'r', encoding="utf-8")  # 读取文件
    yml_config = yaml.load(f, Loader=yaml.FullLoader)  # Loader为了更加安全
    """Loader的几种加载方式
    BaseLoader - -仅加载最基本的YAML
    SafeLoader - -安全地加载YAML语言的子集。建议用于加载不受信任的输入。
    FullLoader - -加载完整的YAML语言。避免任意代码执行。这是当前（PyYAML5.1）默认加载器调用yaml.load(input)（发出警告后）。
 	UnsafeLoader - -（也称为Loader向后兼容性）原始的Loader代码，可以通过不受信任的数据输入轻松利用。"""
    return yml_config


def copy_file(src, dst):
    """
    src 源文件全局路径
    dst 目标路径
    """
    try:
        shutil.copy(src, dst)
    except IOError as e:
        print("Unable to copy file. %s" % e)
    except:
        print("Unexpected error:", sys.exc_info())


def mkdir(path: str):
    '''
    文件夹不存在则新建该文件夹
    '''
    if not os.path.exists(path):
        print("新建目录<{}>".format(path))
        os.makedirs(path, exist_ok=True)


def readcsv(filename, read_header=False):
    '''
    读取csv内的数据
    '''
    csv_datas = dict()
    with open(filename, "r") as csvfile:
        csvreader = csv.reader(csvfile)
        # 遍历csvreader对象的每一行内容并输出
        is_header = True  # 第一行表头不要
        for row in csvreader:
            if not read_header and is_header:
                is_header = False
                continue
            csv_datas[row[0]] = row[1]
    return csv_datas


def read_label(file_path):
    '''
    读取label，返回字典
    '''
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip().split(' ') for line in lines]
    annotation = {}
    if len(lines)==0: # 如果为空，则返回空字典
        annotation['name'] = np.array([])
        annotation['location'] = np.array([])
        annotation['dimensions'] = np.array([])
        annotation['rotation_y'] = np.array([])
        return annotation
    annotation['name'] = np.array([line[0].lower() for line in lines])
    annotation['location'] = np.array([line[1:4] for line in lines],
                                      dtype=np.float32)  # 激光坐标系下的坐标(location) (x, y, z), 下平面中心点的坐标
    annotation['dimensions'] = np.array([line[4:7] for line in lines],
                                        dtype=np.float32)  # 3d 尺寸(dimensions) hwl -> camera coordinates (lhw)
    annotation['rotation_y'] = np.array([line[7:8] for line in lines], dtype=np.float32)  # 相机坐标系下绕轴旋转的弧度(rotation_y)
    return annotation


def read_dr_label(file_path):
    '''
    读取label，返回字典
    '''
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip().split(' ') for line in lines]
    annotation = {}
    annotation['name'] = np.array([line[0].lower() for line in lines])
    annotation['location'] = np.array([line[11:14] for line in lines],
                                      dtype=np.float32)  # 激光坐标系下的坐标(location) (x, y, z), 下平面中心点的坐标
    annotation['dimensions'] = np.array([line[10:7:-1] for line in lines],
                                        dtype=np.float32)  # 3d 尺寸(dimensions) hwl -> camera coordinates (lhw)
    # annotation['dimensions'] = annotation['dimensions'][:, [1, 0, 2]]  # 3d 尺寸(dimensions) hwl -> camera coordinates (lhw)
    annotation['rotation_y'] = - np.array([line[14:15] for line in lines],
                                          dtype=np.float32)  # 相机坐标系下绕轴旋转的弧度(rotation_y)
    annotation['rotation_y'] = annotation['rotation_y'] - np.pi / 2  # 修正旋转角度
    #
    # 预测框坐标系变换：相机坐标系--->lidar坐标系
    tr_velo2cam = np.array([[0.9994015951319486, -0.03199927972993585, 0.013133839670170688, 0.2710357611083472],
                            [0.013650110615193858, 0.015969147292999922, -0.9997793060545548, -0.11226274034223951],
                            [0.03178248146242525, 0.9993603116151348, 0.01639638498544132, -0.02040423686369943],
                            [0, 0, 0, 1]])
    trans_z = np.array([[0.0, 1.0, 0.0, 0.0], [-1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0, 0, 0, 1]])
    tr_cam2velo = tr_velo2cam.copy()
    tr_cam2velo[:3, :3] = np.linalg.inv(tr_velo2cam[:3, :3])
    tr_cam2velo[:3, -1] = np.dot(-tr_cam2velo[:3, :3], tr_velo2cam[:3, -1])

    # 3D框中心点转换坐标--->
    # box_center = box[:, :3]
    box_center1 = np.concatenate((annotation['location'], np.ones((len(annotation['location']), 1))), axis=1)
    annotation['location'] = (trans_z @ tr_cam2velo @ box_center1.T).T[:, :3]
    annotation['location'][:, 2] += 0.5 * annotation['dimensions'][:, 2]  # 高度修正，kitti相机坐标系以地面为中心 lidar坐标系通常以中心点为中心

    return annotation


def read_delft_label(file_path):
    '''
    读取label，返回字典
    '''
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip().split(' ') for line in lines]
    annotation = {}
    annotation['name'] = np.array([line[0].lower() for line in lines])
    annotation['location'] = np.array([line[11:14] for line in lines],
                                      dtype=np.float32)  # 激光坐标系下的坐标(location) (x, y, z), 下平面中心点的坐标
    annotation['dimensions'] = np.array([line[10:7:-1] for line in lines],
                                        dtype=np.float32)  # 3d 尺寸(dimensions) hwl -> camera coordinates (lhw)
    # annotation['dimensions'] = annotation['dimensions'][:, [1, 0, 2]]  # 3d 尺寸(dimensions) hwl -> camera coordinates (lhw)
    annotation['rotation_y'] = - np.array([line[14:15] for line in lines],
                                          dtype=np.float32)  # 相机坐标系下绕轴旋转的弧度(rotation_y)
    annotation['rotation_y'] = annotation['rotation_y'] - np.pi / 2  # 修正旋转角度
    #
    # 预测框坐标系变换：相机坐标系--->lidar坐标系
    tr_velo2cam = np.array(
        [[-0.007980200000000000, -0.999854100000000000, 0.015104900000000000, 0.151000000000000000],
         [0.118497000000000000, -0.015944500000000000, -0.992826400000000000, -0.461000000000000000],
         [0.992922400000000000, -0.006133100000000000, 0.118606900000000000, -0.915000000000000000],
         [0, 0, 0, 1]])
    trans_z = np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0, 0, 0, 1]])
    tr_cam2velo = tr_velo2cam.copy()
    tr_cam2velo[:3, :3] = np.linalg.inv(tr_velo2cam[:3, :3])
    tr_cam2velo[:3, -1] = np.dot(-tr_cam2velo[:3, :3], tr_velo2cam[:3, -1])

    # 3D框中心点转换坐标--->
    # box_center = box[:, :3]
    box_center1 = np.concatenate((annotation['location'], np.ones((len(annotation['location']), 1))), axis=1)
    annotation['location'] = (trans_z @ tr_cam2velo @ box_center1.T).T[:, :3]
    annotation['location'][:, 2] += 0.5 * annotation['dimensions'][:, 2]  # 高度修正，kitti相机坐标系以地面为中心 lidar坐标系通常以中心点为中心

    return annotation


def read_pickle(file_path, suffix='.pkl'):
    assert os.path.splitext(file_path)[1] == suffix
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    return data


def write_pickle(results, file_path):
    with open(file_path, 'wb') as f:
        pickle.dump(results, f)


def read_points(file_path, dim=4, datatype=np.float32):
    suffix = os.path.splitext(file_path)[1]
    assert suffix in ['.bin', '.ply']
    if suffix == '.bin':
        return np.fromfile(file_path, dtype=datatype).reshape(-1, dim)
    else:
        raise NotImplementedError


def write_points(points, file_path, datatype=np.float32):
    points32 = points.astype(datatype)
    suffix = os.path.splitext(file_path)[1]
    assert suffix in ['.bin', '.ply']
    if suffix == '.bin':
        with open(file_path, 'wb') as f:
            points32.tofile(f)
    else:
        raise NotImplementedError


def files_filter(dir_path, pre_fix: list = None, pos_fix: list = None, split='.'):
    files = os.listdir(dir_path)
    files.sort(key=lambda x: x)
    new_files = []
    for file in files:
        file_path = os.path.join(dir_path, file)
        if os.path.isfile(file_path):
            pre = str.split(file, '.')[0]
            pos = str.split(file, '.')[-1]
            pre_ok = True if pre_fix is None else pre in pre_fix
            pos_ok = True if pos_fix is None else pos in pos_fix
            if pre_ok and pos_ok:
                new_files.append(file)
    return new_files
